import numpy as np
from pathlib import Path
import os
import glob
from tqdm import tqdm

def restore_original_npy_labels(labels_npz_path, timestamps_path, output_path):
    """
    从labels.npz和timestamps_us.npy还原原始的.npy标签文件
    
    参数:
        labels_npz_path: labels.npz文件路径
        timestamps_path: timestamps_us.npy文件路径
        output_path: 输出路径
    
    返回:
        还原的标签数组
    """
    # 加载数据
    data = np.load(labels_npz_path)
    labels = data['labels']
    frame_indices = data['objframe_idx_2_label_idx']
    
    timestamps_us = np.load(timestamps_path)
    
    # 重构原始标签
    reconstructed_labels = []
    
    # 遍历每一帧
    for i in range(len(timestamps_us)):
        start_idx = frame_indices[i]
        # 计算结束索引（下一帧的开始或整个标签数组的结束）
        end_idx = frame_indices[i+1] if i+1 < len(frame_indices) else len(labels)
        
        # 获取当前帧的所有标签
        frame_labels = labels[start_idx:end_idx].copy()
        
        # 更新时间戳（如果需要）
        if len(frame_labels) > 0:
            # 确保每个标签都使用帧的时间戳
            frame_labels['t'] = timestamps_us[i]
            reconstructed_labels.append(frame_labels)
    
    # 合并所有标签
    if reconstructed_labels:
        original_labels = np.concatenate(reconstructed_labels)
    else:
        # 如果没有标签，创建一个空数组但保持数据类型
        original_labels = np.array([], dtype=labels.dtype)
    
    # 保存还原的标签
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    np.save(output_path, original_labels)
    
    return original_labels, len(original_labels)

def batch_process_directories(source_dir, target_dir):
    """
    批量处理所有子目录
    
    参数:
        source_dir: 包含所有序列文件夹的源目录
        target_dir: 输出标签文件的目标目录
    """
    # 确保目标目录存在
    os.makedirs(target_dir, exist_ok=True)
    
    # 获取所有子目录
    subdirs = [d for d in Path(source_dir).iterdir() if d.is_dir()]
    print(f"发现 {len(subdirs)} 个子目录")
    
    results = []
    
    # 处理每个子目录
    for subdir in tqdm(subdirs, desc="处理文件夹"):
        # 检查是否存在labels_v2目录
        labels_dir = subdir / "labels_v2"
        if not labels_dir.exists():
            print(f"警告: {subdir.name} 中没有找到labels_v2目录，跳过")
            continue
        
        # 检查必要的文件是否存在
        labels_npz = labels_dir / "labels.npz"
        timestamps_npy = labels_dir / "timestamps_us.npy"
        
        if not labels_npz.exists() or not timestamps_npy.exists():
            print(f"警告: {subdir.name} 中缺少必要的文件，跳过")
            continue
        
        # 构建输出文件路径
        output_path = Path(target_dir) / f"{subdir.name}_bbox.npy"
        
        try:
            # 处理并保存标签
            _, num_labels = restore_original_npy_labels(labels_npz, timestamps_npy, str(output_path))
            results.append((subdir.name, num_labels))
        except Exception as e:
            print(f"处理 {subdir.name} 时出错: {e}")
    
    # 打印处理结果
    print("\n=== 处理结果 ===")
    print(f"总共处理了 {len(results)} 个文件夹")
    for name, count in results:
        print(f"{name}: {count} 个标签")
    
    return results

if __name__ == "__main__":
    # 指定输入和输出目录
    source_directory = r"I:\dataset\\nuclear_power\\nuclear_H5_process\\val"
    target_directory = r"I:\dataset\\lastdata\\H5\\val\\normal"
    
    batch_process_directories(source_directory, target_directory)